#!/usr/bin/env python3
"""
Basic usage examples for VectorDB microservice.

This script demonstrates how to use the VectorDB microservice for basic
vector operations including storing, searching, and managing embeddings.
"""

import asyncio
import httpx
import json
from typing import List, Dict, Any


class VectorDBClient:
    """Simple client for VectorDB microservice."""

    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.client = httpx.AsyncClient(base_url=base_url, timeout=30.0)

    async def health_check(self) -> Dict[str, Any]:
        """Check service health."""
        response = await self.client.get("/api/v1/health")
        response.raise_for_status()
        return response.json()

    async def create_table(self, table_name: str, vector_dimension: int = 1024, overwrite: bool = False) -> Dict[str, Any]:
        """Create a new vector table."""
        data = {
            "table_name": table_name,
            "vector_dimension": vector_dimension,
            "overwrite": overwrite
        }
        response = await self.client.post("/api/v1/tables", json=data)
        response.raise_for_status()
        return response.json()

    async def list_tables(self) -> List[str]:
        """List all tables."""
        response = await self.client.get("/api/v1/tables")
        response.raise_for_status()
        return response.json()

    async def store_vectors(self, table_name: str, texts: List[str],
                           metadata: List[Dict[str, Any]] = None,
                           provider: str = None) -> Dict[str, Any]:
        """Store vectors in a table."""
        data = {
            "table_name": table_name,
            "texts": texts,
            "metadata": metadata or [],
            "provider": provider
        }
        response = await self.client.post("/api/v1/store", json=data)
        response.raise_for_status()
        return response.json()

    async def search_vectors(self, table_name: str, query: str,
                           limit: int = 10, threshold: float = 0.0,
                           provider: str = None) -> Dict[str, Any]:
        """Search for similar vectors."""
        data = {
            "table_name": table_name,
            "query": query,
            "limit": limit,
            "threshold": threshold,
            "provider": provider
        }
        response = await self.client.post("/api/v1/search", json=data)
        response.raise_for_status()
        return response.json()

    async def embed_texts(self, texts: List[str], provider: str = None) -> Dict[str, Any]:
        """Generate embeddings for texts."""
        data = {
            "texts": texts,
            "provider": provider
        }
        response = await self.client.post("/api/v1/embed", json=data)
        response.raise_for_status()
        return response.json()

    async def rag_search(self, query: str, table_name: str = "documents",
                        limit: int = 5, threshold: float = 0.0,
                        provider: str = None) -> Dict[str, Any]:
        """Perform RAG (Retrieval-Augmented Generation) search."""
        data = {
            "query": query,
            "table_name": table_name,
            "limit": limit,
            "threshold": threshold,
            "provider": provider
        }
        response = await self.client.post("/api/v1/rag/search", json=data)
        response.raise_for_status()
        return response.json()

    async def get_database_stats(self) -> Dict[str, Any]:
        """Get database statistics."""
        response = await self.client.get("/api/v1/stats")
        response.raise_for_status()
        return response.json()

    async def close(self):
        """Close the client."""
        await self.client.aclose()


async def basic_example():
    """Basic usage example."""
    print("🚀 VectorDB Basic Usage Example")
    print("=" * 50)

    client = VectorDBClient()

    try:
        # 1. Health check
        print("1. Checking service health...")
        health = await client.health_check()
        print(f"   Service status: {health['status']}")
        print(f"   Database healthy: {health['database']['healthy']}")
        print(f"   Embedding healthy: {health['embedding']['healthy']}")
        print()

        # 2. Create a table
        print("2. Creating table 'documents'...")
        result = await client.create_table("documents", vector_dimension=1024)
        print(f"   Result: {result['message']}")
        print()

        # 3. List tables
        print("3. Listing tables...")
        tables = await client.list_tables()
        print(f"   Available tables: {tables}")
        print()

        # 4. Store some documents
        print("4. Storing documents...")
        documents = [
            "Machine learning is a subset of artificial intelligence.",
            "Deep learning uses neural networks with multiple layers.",
            "Natural language processing enables computers to understand human language.",
            "Computer vision allows machines to interpret and analyze visual information."
        ]

        metadata = [
            {"category": "AI", "topic": "machine_learning"},
            {"category": "AI", "topic": "deep_learning"},
            {"category": "AI", "topic": "nlp"},
            {"category": "AI", "topic": "computer_vision"}
        ]

        result = await client.store_vectors("documents", documents, metadata)
        print(f"   Stored {len(documents)} documents")
        print(f"   Result: {result['message']}")
        print()

        # 5. Search for similar documents
        print("5. Searching for similar documents...")
        query = "What is artificial intelligence?"
        search_results = await client.search_vectors("documents", query, limit=3)

        print(f"   Query: {query}")
        print(f"   Found {search_results['count']} results:")
        for i, result in enumerate(search_results['results'], 1):
            print(f"   {i}. {result['text'][:80]}...")
            print(f"      Score: {result['score']:.4f}")
            print(f"      Metadata: {result['metadata']}")
        print()

        # 6. RAG search
        print("6. Performing RAG search...")
        rag_results = await client.rag_search("Explain machine learning", limit=2)

        print(f"   Query: Explain machine learning")
        print(f"   Retrieved {rag_results['context_count']} context documents:")
        for i, context in enumerate(rag_results['context'], 1):
            print(f"   {i}. {context['text'][:80]}...")
            print(f"      Score: {context['score']:.4f}")
        print()

        # 7. Generate embeddings directly
        print("7. Generating embeddings...")
        texts = ["Hello world", "Vector databases are useful"]
        embeddings = await client.embed_texts(texts)

        print(f"   Generated embeddings for {len(texts)} texts")
        print(f"   Provider: {embeddings['provider']}")
        print(f"   Model: {embeddings['model']}")
        print(f"   Dimensions: {embeddings['dimensions']}")
        print(f"   First embedding (first 5 values): {embeddings['embeddings'][0][:5]}")
        print()

        # 8. Get database stats
        print("8. Database statistics...")
        stats = await client.get_database_stats()
        print(f"   Database URI: {stats['uri']}")
        print(f"   Total tables: {stats['total_tables']}")
        print(f"   Tables: {list(stats['tables'].keys())}")
        print(f"   Connection pool stats: {stats['pool_stats']}")
        print()

        print("✅ Basic example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await client.close()


async def advanced_example():
    """Advanced usage example with multiple tables and filtering."""
    print("🚀 VectorDB Advanced Usage Example")
    print("=" * 50)

    client = VectorDBClient()

    try:
        # Create multiple tables for different content types
        print("1. Creating multiple specialized tables...")

        # Academic papers table
        await client.create_table("academic_papers", vector_dimension=1024)

        # Product reviews table
        await client.create_table("product_reviews", vector_dimension=1024)

        # FAQ table
        await client.create_table("faq", vector_dimension=1024)

        tables = await client.list_tables()
        print(f"   Created tables: {[t for t in tables if t in ['academic_papers', 'product_reviews', 'faq']]}")
        print()

        # Store academic papers
        print("2. Storing academic papers...")
        papers = [
            "Attention Is All You Need: We propose a new simple network architecture, the Transformer.",
            "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.",
            "GPT-3: Language Models are Few-Shot Learners.",
            "DALL-E: Creating Images from Text with Contrastive Learning."
        ]

        paper_metadata = [
            {"type": "paper", "year": 2017, "conference": "NIPS", "authors": ["Vaswani", "Shazeer", "Parmar"]},
            {"type": "paper", "year": 2019, "conference": "NAACL", "authors": ["Devlin", "Chang", "Lee"]},
            {"type": "paper", "year": 2020, "conference": "NeurIPS", "authors": ["Brown", "Mann", "Ryder"]},
            {"type": "paper", "year": 2021, "conference": "ICML", "authors": ["Ramesh", "Pavlov", "Goh"]}
        ]

        result = await client.store_vectors("academic_papers", papers, paper_metadata)
        print(f"   Stored {len(papers)} academic papers")
        print()

        # Store product reviews
        print("3. Storing product reviews...")
        reviews = [
            "This laptop is amazing! Fast performance and great battery life.",
            "The screen quality is excellent, very sharp and vibrant colors.",
            "Keyboard feels comfortable for long typing sessions.",
            "Build quality is solid, feels very durable.",
            "Price is a bit high but worth it for the quality."
        ]

        review_metadata = [
            {"product": "laptop", "rating": 5, "category": "electronics"},
            {"product": "laptop", "rating": 5, "category": "display"},
            {"product": "laptop", "rating": 4, "category": "keyboard"},
            {"product": "laptop", "rating": 5, "category": "build_quality"},
            {"product": "laptop", "rating": 4, "category": "value"}
        ]

        result = await client.store_vectors("product_reviews", reviews, review_metadata)
        print(f"   Stored {len(reviews)} product reviews")
        print()

        # Cross-table search
        print("4. Cross-table semantic search...")
        queries = [
            "transformer architecture",
            "computer performance",
            "display quality"
        ]

        for query in queries:
            print(f"   Query: '{query}'")

            # Search in academic papers
            paper_results = await client.search_vectors("academic_papers", query, limit=2)
            if paper_results['count'] > 0:
                print(f"   📚 Academic Papers: {paper_results['results'][0]['text'][:60]}...")

            # Search in product reviews
            review_results = await client.search_vectors("product_reviews", query, limit=2)
            if review_results['count'] > 0:
                print(f"   💻 Product Reviews: {review_results['results'][0]['text'][:60]}...")
            print()

        # Advanced RAG with context
        print("5. Advanced RAG with multi-source context...")

        complex_query = "What are the latest developments in AI and machine learning?"

        # Get context from academic papers
        paper_context = await client.rag_search(complex_query, "academic_papers", limit=3)

        # Get context from other sources (simulated)
        print(f"   Query: {complex_query}")
        print(f"   Retrieved {paper_context['context_count']} academic papers:")
        for i, ctx in enumerate(paper_context['context'], 1):
            print(f"   {i}. [{ctx['metadata'].get('year', 'N/A')}] {ctx['text'][:80]}...")
        print()

        print("✅ Advanced example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await client.close()


async def batch_operations_example():
    """Example of batch operations for efficiency."""
    print("🚀 VectorDB Batch Operations Example")
    print("=" * 50)

    client = VectorDBClient()

    try:
        # Prepare large dataset
        print("1. Preparing large dataset...")

        # Generate synthetic documents
        documents = []
        topics = ["technology", "science", "health", "education", "business"]

        for topic in topics:
            for i in range(20):  # 20 documents per topic
                documents.append({
                    "text": f"{topic.capitalize()} document {i}: This is a comprehensive article about {topic} innovations and trends in 2024.",
                    "metadata": {
                        "topic": topic,
                        "doc_id": f"{topic}_{i}",
                        "timestamp": f"2024-{i%12+1:02d}-01"
                    }
                })

        print(f"   Generated {len(documents)} documents across {len(topics)} topics")
        print()

        # Batch store documents
        print("2. Batch storing documents...")

        batch_size = 50
        total_stored = 0

        for i in range(0, len(documents), batch_size):
            batch = documents[i:i + batch_size]
            texts = [doc["text"] for doc in batch]
            metadata = [doc["metadata"] for doc in batch]

            result = await client.store_vectors("documents", texts, metadata)
            total_stored += len(batch)
            print(f"   Stored batch {i//batch_size + 1}: {len(batch)} documents")

        print(f"   Total stored: {total_stored} documents")
        print()

        # Batch search operations
        print("3. Batch search operations...")

        search_queries = [
            "technology trends 2024",
            "scientific discoveries",
            "healthcare innovations",
            "educational technology",
            "business strategies"
        ]

        search_tasks = []
        for query in search_queries:
            task = client.search_vectors("documents", query, limit=5)
            search_tasks.append(task)

        # Execute all searches concurrently
        results = await asyncio.gather(*search_tasks)

        for query, result in zip(search_queries, results):
            print(f"   Query: '{query}'")
            print(f"   Found {result['count']} results")
            if result['results']:
                print(f"   Top result: {result['results'][0]['text'][:60]}...")
            print()

        # Performance comparison
        print("4. Performance metrics...")

        import time

        # Single search
        start_time = time.time()
        single_result = await client.search_vectors("documents", "technology", limit=10)
        single_time = time.time() - start_time

        # Batch search (5 queries)
        start_time = time.time()
        batch_queries = ["technology", "science", "health", "education", "business"]
        batch_tasks = [client.search_vectors("documents", q, limit=5) for q in batch_queries]
        batch_results = await asyncio.gather(*batch_tasks)
        batch_time = time.time() - start_time

        print(f"   Single search time: {single_time:.3f}s")
        print(f"   Batch search time (5 queries): {batch_time:.3f}s")
        print(f"   Average per query: {batch_time/5:.3f}s")
        print()

        print("✅ Batch operations example completed successfully!")

    except Exception as e:
        print(f"❌ Error: {e}")
        raise

    finally:
        await client.close()


async def main():
    """Run all examples."""
    print("VectorDB Microservice Examples")
    print("=" * 60)
    print()

    # Run basic example
    await basic_example()
    print()
    print("-" * 60)
    print()

    # Run advanced example
    await advanced_example()
    print()
    print("-" * 60)
    print()

    # Run batch operations example
    await batch_operations_example()


if __name__ == "__main__":
    asyncio.run(main())